#ifndef _H_S_TASK_H_
#define _H_S_TASK_H_

#include "../typedefine.h"
#include <coroutine>
#include <memory>
#include <functional>
#include "../threadpool.h"
#include "../promise.h"

namespace TB
{
	template<typename T> struct STaskProcessor;

	template<typename T = void>
	struct STask
	{
	private:
		friend struct STaskProcessor<T>;
		friend class CTaskHelper;
		class CInfo
		{
		public:
			CInfo():
				m_oRet()
			{
			}

			bool Resolve(const T& oValue)
			{
				if (m_bIsCompleted)
					return false;
				std::vector<std::function<void()>> oNextTasks;
				{
					CSpinLock oLock(m_oSpinLockFlag);
					if (m_bIsCompleted)
						return false;
					m_oRet = oValue;
					m_bIsCompleted = true;
					oNextTasks.swap(m_oNextTasks);
				}
				for (UINT32 i = 0; i != oNextTasks.size(); ++i)
				{
					oNextTasks[i]();
				}
				return true;
			}

			bool Resolve(T&& oValue)
			{
				if (m_bIsCompleted)
					return false;
				std::vector<std::function<void()>> oNextTasks;
				{
					CSpinLock oLock(m_oSpinLockFlag);
					if (m_bIsCompleted)
						return false;
					m_oRet = std::forward<T>(oValue);
					m_bIsCompleted = true;
					oNextTasks.swap(m_oNextTasks);
				}
				for (UINT32 i = 0; i != oNextTasks.size(); ++i)
				{
					oNextTasks[i]();
				}
				return true;
			}

			bool Then(const std::function<void()>& rcbTask)
			{
				if (!m_bIsCompleted)
				{
					CSpinLock oLock(m_oSpinLockFlag);
					if (!m_bIsCompleted)
					{
						m_oNextTasks.push_back(rcbTask);
						return true;
					}
				}
				return false;
			}
		public:
			bool m_bIsCompleted = false;
			T m_oRet;
			std::vector<std::function<void()>> m_oNextTasks;
			CSpinLockFlag m_oSpinLockFlag;
		};

	public:
		struct promise_type {
			STask<T> get_return_object() 
			{
				STask<T> sTask;
				m_poInfo = sTask.m_poInfo;
				return sTask;
			}
			std::suspend_never initial_suspend() 
			{
				return {};
			}
			std::suspend_never final_suspend() noexcept 
			{
				return {};
			}
			void return_value(T&& oValue) 
			{
				m_poInfo->m_oRet = std::forward<T>(oValue);
			}
			void unhandled_exception()
			{
				throw;
			}
			std::shared_ptr<CInfo> m_poInfo;
		};

	public: // awaiter
		bool await_ready() 
		{
			return m_poInfo == nullptr || m_poInfo->m_bIsCompleted;
		}
		bool await_suspend(std::coroutine_handle<> handle) 
		{
			CCurThreadSyncContext& roCurThreadSyncContext = CCurThreadSyncContext::GetCurrent();
			return m_poInfo->Then([handle, &roCurThreadSyncContext]()
			{
				auto poContext = roCurThreadSyncContext.GetSynchronizationContext();
				if (poContext == nullptr)
				{
					handle();
					return;
				}
				poContext->Post(handle);
			});
		}
		T await_resume() 
		{
			return m_poInfo == nullptr ? m_oValue : m_poInfo->m_oRet;
		}
	private:
		STask() : m_oValue(), m_poInfo(std::make_shared<CInfo>()) {}
		explicit STask(const T& oValue) : m_oValue(oValue) {}
		explicit STask(T&& oValue) : m_oValue(oValue) {}

	public:
		STask(const STask& oTask) noexcept :m_oValue(oTask.m_oValue), m_poInfo(oTask.m_poInfo) {}
		STask& operator=(const STask& oTask) noexcept { m_oValue = oTask.m_oValue; m_poInfo = oTask.m_poInfo; return *this; }
		STask(STask&& oTask) noexcept :m_oValue(std::move(oTask.m_oValue)), m_poInfo(oTask.m_poInfo) {}
		STask& operator=(STask&& oTask) noexcept { m_oValue = std::move(oTask.m_oValue); m_poInfo = oTask.m_poInfo; return *this; }

	public:
		inline bool IsCompleted() const { return m_poInfo == nullptr || m_poInfo->m_bIsCompleted; }
		inline const T& GetResult() const { return m_poInfo == nullptr ? m_oValue : m_poInfo->m_oRet; }

	private:
		T m_oValue;
		std::shared_ptr<CInfo> m_poInfo;
	};

	template<>
	struct STask<void>
	{
	private:
		friend struct STaskProcessor<void>;
		friend class CTaskHelper;
		class CInfo
		{
		public:
			CInfo() :
				m_bIsCompleted(false)
			{
			}

			bool Resolve()
			{
				if (m_bIsCompleted)
					return false;
				std::vector<std::function<void()>> oNextTasks;
				{
					CSpinLock oLock(m_oSpinLockFlag);
					if (m_bIsCompleted)
						return false;
					m_bIsCompleted = true;
					oNextTasks.swap(m_oNextTasks);
				}
				for (UINT32 i = 0; i != oNextTasks.size(); ++i)
				{
					oNextTasks[i]();
				}
				return true;
			}

			bool Then(const std::function<void()>& rcbTask)
			{
				if (!m_bIsCompleted)
				{
					CSpinLock oLock(m_oSpinLockFlag);
					if (!m_bIsCompleted)
					{
						m_oNextTasks.push_back(rcbTask);
						return true;
					}
				}
				return false;
			}
		public:
			bool m_bIsCompleted = false;
			std::vector<std::function<void()>> m_oNextTasks;
			CSpinLockFlag m_oSpinLockFlag;
		};

	public:
		struct promise_type {
			STask<void> get_return_object()
			{
				return STask<void>(true);
			}
			std::suspend_never initial_suspend()
			{
				return {};
			}
			std::suspend_never final_suspend() noexcept
			{
				return {};
			}
			void return_void()
			{
			}
			void unhandled_exception()
			{
				throw;
			}
		};

	public: // awaiter
		bool await_ready()
		{
			return m_poInfo == nullptr || m_poInfo->m_bIsCompleted;
		}
		bool await_suspend(std::coroutine_handle<> handle)
		{
			CCurThreadSyncContext& roCurThreadSyncContext = CCurThreadSyncContext::GetCurrent();
			return m_poInfo->Then([handle, &roCurThreadSyncContext]()
			{
				auto poContext = roCurThreadSyncContext.GetSynchronizationContext();
				if (poContext == nullptr)
				{
					handle();
					return;
				}
				poContext->Post(handle);
			});
		}
		void await_resume()
		{
		}

	private:
		explicit STask(bool bCompleted) 
		{
			if (!bCompleted)
				m_poInfo = std::make_shared<CInfo>();
		}
	public:
		STask(const STask& oTask) noexcept : m_poInfo(oTask.m_poInfo) {}
		STask& operator=(const STask& oTask) noexcept { m_poInfo = oTask.m_poInfo; return *this; }
		STask(STask&& oTask) noexcept : m_poInfo(oTask.m_poInfo) {}
		STask& operator=(STask&& oTask) noexcept { m_poInfo = oTask.m_poInfo; return *this; }

	public:
		inline bool IsCompleted() const { return m_poInfo == nullptr || m_poInfo->m_bIsCompleted; }

	private:
		std::shared_ptr<CInfo> m_poInfo;
	};
}

#endif
